import copy
import json
from tracemalloc import start
import cv2
from cgi import test
from cv2 import mean
from importlib_metadata import distribution, re
import numpy as np
import math
import time
import sys
import os
import argparse
from colorama import Fore, Back, Style
from colorama import init
import matplotlib.pyplot as plt
from scipy.special import entr
from scipy import stats
from collections import Counter

import threading
import multiprocessing

from c_acc import propagate2
from c_acc import update_wave_entropy_c

###############################################################################
# 2D wave function collapse
# input:
#   rules: a list of rules
###############################################################################

Width = 9
Height = 9
entropy = []
# 1-5 bricks
# 6-9 corners
# 10-13 ramps
candidates = [1,2,3,4,5,6,7,8,9,10,11,12,13]
candidates_X_symmetry = {1,2,3,4,5}
candidates_L_symmetry = {6,7,8,9,10,11,12,13}
candidates_weight = [1,1,1,1,1,1,1,1,1,1,1,1,1]
weight_stack = []

oriented_candidates = []
oriented_candidates_weight = []
oriented_weight_stack = []

# weight dictionary for each cell
oriented_weight_dict = []
oriented_candidates_weight_dict = {}

# the wave
# wave is a 2D array of size (W x H, number_of_candidates)
wave_oriented = []

# fast entropy calculation
entropy_lut = []
def build_entropy_lut():
    for i in range(1,10):
        entropy_lut.append((-(i*0.1) * math.log(i*0.1)))
def fast_entropy(p):
    return entropy_lut[(int)(p*9)]

# all the compatible situations
# H for horizontal
# V for vertical
# Rotation: 0->0 , 1->90, 2->180, 3->270

rules = {
    (1,0,2,0,'H'),
    (2,0,3,0,'H'),
    (3,0,4,0,'H'),
    (4,0,5,0,'H'),
    (5,0,4,0,'H'),
    (4,0,3,0,'H'),
    (3,0,2,0,'H'),
    (2,0,1,0,'H'),
    (1,0,1,0,'H'),
    (2,0,2,0,'H'),
    (3,0,3,0,'H'),
    (4,0,4,0,'H'),
    (5,0,5,0,'H'),

    (6,0,1,0,'H'),
    (2,0,6,0,'H'),
    (1,0,6,0,'V'),
    (6,0,2,0,'V'),

    (7,0,2,0,'H'),
    (3,0,7,0,'H'),
    (2,0,7,0,'V'),
    (7,0,3,0,'V'),

    (8,0,3,0,'H'),
    (4,0,8,0,'H'),
    (3,0,8,0,'V'),
    (8,0,4,0,'V'),

    (9,0,4,0,'H'),
    (5,0,9,0,'H'),
    (4,0,9,0,'V'),
    (9,0,5,0,'V'),

    (2,0,10,0,'V'),
    (10,0,1,0,'V'),
    (1,0,10,0,'H'),
    (10,0,10,0,'H'),
    (10,0,1,0,'H'),

    (3,0,11,0,'V'),
    (11,0,2,0,'V'),
    (2,0,11,0,'H'),
    (11,0,11,0,'H'),
    (11,0,2,0,'H'),

    (4,0,12,0,'V'),
    (12,0,3,0,'V'),
    (3,0,12,0,'H'),
    (12,0,12,0,'H'),
    (12,0,3,0,'H'),

    (5,0,13,0,'V'),
    (13,0,4,0,'V'),
    (4,0,13,0,'H'),
    (13,0,13,0,'H'),
    (13,0,4,0,'H')
}

rules_h = set()
rules_v = set()

def rotate_clockwise(tp = (0,0)):

    if tp[0] in candidates_X_symmetry:
        return tp
    else:
        if tp[1] <=2:
            tp1 = (tp[0],tp[1]+1)
            return tp1
        elif tp[1] == 3:
            tp1 = (tp[0],0)
            return tp1
        else:
            return tp

# symmetry generator
def symmetry_generator(rules2):

    print(rules2)

    rules_ = set()

    # for each rule in rules2, rotate clockwise 4 times
    for r in rules2:

        cand1 = (r[0],r[1])
        cand2 = (r[2],r[3])

        # rot 0
        rule_0 = r

        # rot 90
        rot_0_1 = rotate_clockwise(cand1)
        rot_0_2 = rotate_clockwise(cand2)
        if r[4] == 'H':
            rule_90 = (rot_0_1[0],rot_0_1[1],rot_0_2[0],rot_0_2[1],'V')
        else:
            rule_90 = (rot_0_2[0],rot_0_2[1],rot_0_1[0],rot_0_1[1],'H')

        # rot 180
        rot_90_1 = rotate_clockwise(rotate_clockwise(cand1))
        rot_90_2 = rotate_clockwise(rotate_clockwise(cand2))
        if r[4] == 'H':
            rule_180 = (rot_90_2[0],rot_90_2[1],rot_90_1[0],rot_90_1[1],'H')
        else:
            rule_180 = (rot_90_2[0],rot_90_2[1],rot_90_1[0],rot_90_1[1],'V')

        # rot 270
        rot_270_1 = rotate_clockwise(rotate_clockwise(rotate_clockwise(cand1)))
        rot_270_2 = rotate_clockwise(rotate_clockwise(rotate_clockwise(cand2)))
        if r[4] == 'H':
            rule_270 = (rot_270_2[0],rot_270_2[1],rot_270_1[0],rot_270_1[1],'V')
        else:
            rule_270 = (rot_270_1[0],rot_270_1[1],rot_270_2[0],rot_270_2[1],'H')

        # add to rules
        rules_.add(rule_0)
        rules_.add(rule_90)
        rules_.add(rule_180)
        rules_.add(rule_270)

    print(rules_)
    return rules_

# separate rules into horizontal and vertical
def separate_rules(rules):
    rules_h = set()
    rules_v = set()
    

    for r in rules:
        if r[4] == 'H':
            rules_h.add((r[0],r[1],r[2],r[3]))
        else:
            rules_v.add((r[0],r[1],r[2],r[3]))

    return rules_h,rules_v

# collapse state detection
def collapse_detection(wave_index):
    if len(wave_oriented[wave_index]) != 1:
        return False
    else:
        return True

def distributional_trigger(wavelet_index):

    collapse_element_index = 0

    # construct a new weight array for the wavelet
    weight_array = []
    for i in range(len(wave_oriented[wavelet_index])):
        weight_array.append(get_weight_of_candidate2(wavelet_index,wave_oriented[wavelet_index][i]))

    # generate a random number in [0, sum(weight_array)]
    random_number = np.random.rand() * sum(weight_array)

    # find the collapse element
    for i in range(len(weight_array)):
        random_number -= weight_array[i]
        if random_number <= 0:
            collapse_element_index = i
            break

    return collapse_element_index

# get weight through candidate
def get_weight_of_candidate(index, candidate):
    for i in range(len(oriented_candidates)):
        if oriented_candidates[i] == candidate:
            return oriented_weight_stack[index][i]

def get_weight_of_candidate2(index, candidate):
    return oriented_weight_dict[index][candidate]

def set_weight_of_candidate(index, candidate, weight):
    for i in range(len(candidates)):
        if candidates[i] == candidate:
            weight_stack[index][i] = weight

def Entropy(labels):
    prob_dict = {x:labels.count(x)/len(labels) for x in labels}
    probs = np.array(list(prob_dict.values()))

    return - probs.dot(np.log2(probs))

# calculate the Shannon entropy of a wavelet
def shannon_entropy_of_wavelet(index,wavelet):

    entropy = 0
    sum_weight = 0
    probabilities = []
    for i in wavelet:
        probabilities.append(get_weight_of_candidate2(index,i))

    sum_weight = sum(probabilities)
    probabilities = [x/sum_weight for x in probabilities]
    probabilities_arr = np.array(probabilities)

    entropy = -probabilities_arr.dot(np.log(probabilities))

    entropy += np.random.rand() * 0.001

    return entropy

# update entropy
def update_wave_entropy():
    for i in range(len(wave_oriented)):
        entropy[i] = shannon_entropy_of_wavelet(i,wave_oriented[i])

# find an element to collapse
def find_minimum_entropy_position():
    # find minimum entropy index of entropy
    min_entropy_index = 0
    min_entropy = 1e10
    for i in range(len(entropy)):
        if entropy[i] < min_entropy and collapse_detection(i) == False:
            min_entropy = entropy[i]
            min_entropy_index = i

    return min_entropy_index

# propagate the wave function
# active-propagation
def propagate():
    
    # neighbors :
    #         (i,j-1)
    # (i-1,j) (i,j  ) (i+1,j)
    #         (i,j+1)

    for k in range(1):
        for j in range(Height):
            for i in range(Width):

                # if len(wave[j*Width+i]) == len(candidates):
                #     continue

                # if collapse_detection(j*Width+i) == True:
                #     continue

                i_0 = i-1
                j_0 = j-1
                i_1 = i+1
                j_1 = j+1

                wavelet_temp1 = set()
                wavelet_temp2 = set()
                wavelet_temp3 = set()
                wavelet_temp4 = set()

                # construct local context and look it up in rules
                # (i-1,j)
                if i_0 >= 0:
                    for element_i in wave_oriented[j*Width+i]:
                        for element_j in wave_oriented[j*Width+i_0]:
                            if len(wave_oriented[j*Width+i_0]) == 1:
                                break
                            if (element_j[0],element_j[1],element_i[0],element_i[1],'H') in rules:
                                wavelet_temp1.add(element_j)
                # (i+1,j)
                if i_1 < Width:
                    for element_i in wave_oriented[j*Width+i]:
                        for element_j in wave_oriented[j*Width+i_1]:
                            if len(wave_oriented[j*Width+i_1]) == 1:
                                break
                            if (element_i[0],element_i[1],element_j[0],element_j[1],'H') in rules:
                                wavelet_temp2.add(element_j)
                # (i,j-1)
                if j_0 >= 0:
                    for element_i in wave_oriented[j*Width+i]:
                        for element_j in wave_oriented[j_0*Width+i]:
                            if len(wave_oriented[j_0*Width+i]) == 1:
                                break
                            if (element_j[0],element_j[1],element_i[0],element_i[1],'V') in rules:
                                wavelet_temp3.add(element_j)
                # (i,j+1)
                if j_1 < Height:
                    for element_i in wave_oriented[j*Width+i]:
                        for element_j in wave_oriented[j_1*Width+i]:
                            if len(wave_oriented[j_1*Width+i]) == 1:
                                break
                            if (element_i[0],element_i[1],element_j[0],element_j[1],'V') in rules:
                                wavelet_temp4.add(element_j)
                
                # update neighbors
                if len(wavelet_temp1) != 0:
                    wave_oriented[j*Width+i_0].clear()
                    for element in wavelet_temp1:
                        wave_oriented[j*Width+i_0].append(element)
                if len(wavelet_temp2) != 0:
                    wave_oriented[j*Width+i_1].clear()
                    for element in wavelet_temp2:
                        wave_oriented[j*Width+i_1].append(element)
                if len(wavelet_temp3) != 0:
                    wave_oriented[j_0*Width+i].clear()
                    for element in wavelet_temp3:
                        wave_oriented[j_0*Width+i].append(element)
                if len(wavelet_temp4) != 0:
                    wave_oriented[j_1*Width+i].clear()
                    for element in wavelet_temp4:
                        wave_oriented[j_1*Width+i].append(element)

    update_wave_entropy()

# fake propagation for time-consumption measurement
def fake_propagate():
    
    # neighbors :
    #         (i,j-1)
    # (i-1,j) (i,j  ) (i+1,j)
    #         (i,j+1)

    p1 = 0
    p2 = 0
    p3 = 0
    p4 = 0

    for k in range(1):
        for j in range(Height):
            for i in range(Width):

                # if len(wave[j*Width+i]) == len(candidates):
                #     continue

                # if collapse_detection(j*Width+i) == True:
                #     continue

                i_0 = i-1
                j_0 = j-1
                i_1 = i+1
                j_1 = j+1

                # construct local context and look it up in rules
                w1 = wave_oriented[j*Width+i].copy()
                # (i-1,j)
                if i_0 >= 0 and len(wave_oriented[j*Width+i_0]) != 1:
                    w2 = wave_oriented[j*Width+i_0].copy()
                    for element_i in w1:
                        for element_j in w2:
                            if (element_j[0],element_j[1],element_i[0],element_i[1],'H') in rules_h:
                                p1 += 1
                # (i+1,j)
                if i_1 < Width and len(wave_oriented[j*Width+i_1]) != 1:
                    w2 = wave_oriented[j*Width+i_1].copy()
                    for element_i in w1:
                        for element_j in w2:
                            if (element_i[0],element_i[1],element_j[0],element_j[1],'H') in rules_h:
                                p2 += 1
                # (i,j-1)
                if j_0 >= 0 and len(wave_oriented[j_0*Width+i]) != 1:
                    w2 = wave_oriented[j_0*Width+i].copy()
                    for element_i in w1:
                        for element_j in w2:
                            if (element_j[0],element_j[1],element_i[0],element_i[1],'V') in rules_v:
                                p3 += 1
                # (i,j+1)
                if j_1 < Height and len(wave_oriented[j_1*Width+i]) != 1:
                    w2 = wave_oriented[j_1*Width+i].copy()
                    for element_i in w1:
                        for element_j in w2:
                            if (element_i[0],element_i[1],element_j[0],element_j[1],'V') in rules_v:
                                p4 += 1
                
                # # update neighbors
                # if len(wavelet_temp1) != 0:
                #     wave_oriented[j*Width+i_0].clear()
                #     for element in wavelet_temp1:
                #         wave_oriented[j*Width+i_0].append(element)
                # if len(wavelet_temp2) != 0:
                #     wave_oriented[j*Width+i_1].clear()
                #     for element in wavelet_temp2:
                #         wave_oriented[j*Width+i_1].append(element)
                # if len(wavelet_temp3) != 0:
                #     wave_oriented[j_0*Width+i].clear()
                #     for element in wavelet_temp3:
                #         wave_oriented[j_0*Width+i].append(element)
                # if len(wavelet_temp4) != 0:
                #     wave_oriented[j_1*Width+i].clear()
                #     for element in wavelet_temp4:
                #         wave_oriented[j_1*Width+i].append(element)
    
    # print(p1,p2,p3,p4)


def fake_propagate1():
    
    # neighbors :
    #         (i,j-1)
    # (i-1,j) (i,j  ) (i+1,j)
    #         (i,j+1)
    p1 = 0
    for k in range(1):
        for j in range(Height):
            for i in range(Width):
                i_0 = i-1
                # construct local context and look it up in rules
                w1 = wave_oriented[j*Width+i].copy()
                # (i-1,j)
                if i_0 >= 0 and len(wave_oriented[j*Width+i_0]) != 1:
                    w2 = wave_oriented[j*Width+i_0].copy()
                    ll = set()
                    for element_i in w1:
                        for element_j in w2:
                            if (element_j[0],element_j[1],element_i[0],element_i[1],'H') in rules_h:
                                p1 += 1
    print(p1)

# collapse the wavelet
def collapse_at_position(wavelet_index):
    
    collapse_element_index = distributional_trigger(wavelet_index)
    wavelet = []
    wavelet.append(wave_oriented[wavelet_index][collapse_element_index])
    wave_oriented[wavelet_index].clear()
    wave_oriented[wavelet_index] = wavelet.copy()

def clear_screen():
    print('\033c',end='')

# check if the wave is collapsed
def is_collapsed(wave_oriented):
    for i in range(len(wave_oriented)):
        if collapse_detection(i) == False:
            return False
    return True

# contradiction check
def contradiction_check():
    for j in range(Height):
        for i in range(Width):

            if collapse_detection(j*Width+i) == False:
                print(j*Width+i)
                continue

            i_0 = i-1
            j_0 = j-1
            i_1 = i+1
            j_1 = j+1

            # construct local context and compare it with rules
            # (i-1,j)
            if i_0 >= 0:
                for element_i in wave_oriented[j*Width+i]:
                    for element_j in wave_oriented[j*Width+i_0]:
                        if not (element_j[0],element_j[1],element_i[0],element_i[1],'H') in rules:
                            return False
            # (i+1,j)
            if i_1 < Width:
                for element_i in wave_oriented[j*Width+i]:
                    for element_j in wave_oriented[j*Width+i_1]:
                        if not (element_i[0],element_i[1],element_j[0],element_j[1],'H') in rules:
                            return False
            # (i,j-1)
            if j_0 >= 0:
                for element_i in wave_oriented[j*Width+i]:
                    for element_j in wave_oriented[j_0*Width+i]:
                        if not (element_j[0],element_j[1],element_i[0],element_i[1],'V') in rules:
                            return False
            # (i,j+1)
            if j_1 < Height:
                for element_i in wave_oriented[j*Width+i]:
                    for element_j in wave_oriented[j_1*Width+i]:
                        if not (element_i[0],element_i[1],element_j[0],element_j[1],'V') in rules:
                            return False
    return True

def render(wave_oriented, wind_name = "canvas",write_ = False,write_id = 0,output = False):

    # bricks
    tile1 = np.zeros((20,20,3), np.uint8)
    tile2 = np.zeros((20,20,3), np.uint8)
    tile2[:,:] = (50,50,50)
    tile3 = np.zeros((20,20,3), np.uint8)
    tile3[:,:] = (100,100,100)
    tile4 = np.zeros((20,20,3), np.uint8)
    tile4[:,:] = (150,150,150)
    tile5 = np.zeros((20,20,3), np.uint8)
    tile5[:,:] = (200,200,200)

    # corners
    tile6 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            if j<i:
                tile6[j,i] = (0,0,0)
            else:
                tile6[j,i] = (50,50,50)
    tile7 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            if j<i:
                tile7[j,i] = (50,50,50)
            else:
                tile7[j,i] = (100,100,100)
    tile8 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            if j<i:
                tile8[j,i] = (100,100,100)
            else:
                tile8[j,i] = (150,150,150)
    tile9 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            if j<i:
                tile9[j,i] = (150,150,150)
            else:
                tile9[j,i] = (200,200,200)

    # ramps
    tile10 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            tile10[j,i] = (50-j*2.5,50-j*2.5,50-j*2.5)
    tile11 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            tile11[j,i] = (100-j*2.5,100-j*2.5,100-j*2.5)
    tile12 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            tile12[j,i] = (150-j*2.5,150-j*2.5,150-j*2.5)
    tile13 = np.zeros((20,20,3), np.uint8)
    for j in range(20):
        for i in range(20):
            tile13[j,i] = (200-j*2.5,200-j*2.5,200-j*2.5)

    tiles = [] 
    tiles.append(tile1)
    tiles.append(tile2)
    tiles.append(tile3)
    tiles.append(tile4)
    tiles.append(tile5)
    tiles.append(tile6)
    tiles.append(tile7)
    tiles.append(tile8)
    tiles.append(tile9)
    tiles.append(tile10)
    tiles.append(tile11)
    tiles.append(tile12)
    tiles.append(tile13)
    # draw a black rectangle on the border of each tile
    for i in range(len(tiles)):
        tiles[i][0,:] = 0
        tiles[i][:,0] = 0
        tiles[i][19,:] = 0
        tiles[i][:,19] = 0

    # conbine Height x Width tiles into a single image
    canvas = np.zeros((Height*20,Width*20,3), np.uint8)

    for j in range(Height):
        for i in range(Width):
            if collapse_detection(j*Width+i) == True:
                if wave_oriented[j*Width+i][0][1] == 1:
                    # rotate the tile 90 degrees clockwise
                    canvas[j*20:(j+1)*20,i*20:(i+1)*20] = np.rot90(tiles[wave_oriented[j*Width+i][0][0]-1],-1).copy()
                elif wave_oriented[j*Width+i][0][1] == 2:
                    # rotate the tile 180 degrees clockwise
                    canvas[j*20:(j+1)*20,i*20:(i+1)*20] = np.rot90(tiles[wave_oriented[j*Width+i][0][0]-1],-2).copy()
                elif wave_oriented[j*Width+i][0][1] == 3:
                    # rotate the tile 270 degrees clockwise
                    canvas[j*20:(j+1)*20,i*20:(i+1)*20] = np.rot90(tiles[wave_oriented[j*Width+i][0][0]-1],-3).copy()
                else:
                    canvas[j*20:(j+1)*20,i*20:(i+1)*20] = tiles[wave_oriented[j*Width+i][0][0]-1].copy()

    if not output:
        for j in range(Height):
            for i in range(Width):
                sum_ = len(wave_oriented[j*Width+i])==1 and wave_oriented[j*Width+i][0][0] or len(wave_oriented[j*Width+i])
                dir_ = wave_oriented[j*Width+i][0][1]
                mean = sum_
                mean_str = str(round(mean,2)) + ":" + str(dir_)
                if len(wave_oriented[j*Width+i]) == len(oriented_candidates):
                    print(Style.RESET_ALL + Fore.CYAN + mean_str + "," + Fore.LIGHTBLUE_EX, end=' ')
                elif len(wave_oriented[j*Width+i]) == 1:
                    print(Style.RESET_ALL + Fore.GREEN + mean_str + ",", end=' ')
                else:
                    print(Style.RESET_ALL + Fore.RED  + mean_str + ",", end=' ')
            print()

    if write_ == True:
        id_string = str(write_id)
        cv2.imwrite(id_string + ".png", canvas)
        # save wave to a json file
        with open(id_string + ".json", 'w') as outfile:
            json.dump(wave_oriented, outfile)
        
    if not output:
        cv2.imshow(wind_name,canvas)
        cv2.waitKey(1)
    else:
        return canvas

def super_render(wave1,wave2,wave3):

    cav1 = render(wave1,output=True)
    cav2 = render(wave2,output=True)
    cav3 = render(wave3,output=True)

    # make a green border for all three images
    
    cav1[0,:] = (0,255,0)
    cav1[:,0] = (0,255,0)
    cav1[Height*20-1,:] = (0,255,0)
    cav1[:,Width*20-1] = (0,255,0)

    cav2[0,:] = (0,255,0)
    cav2[:,0] = (0,255,0)
    cav2[Height*20-1,:] = (0,255,0)
    cav2[:,Width*20-1] = (0,255,0)

    cav3[0,:] = (0,255,0)
    cav3[:,0] = (0,255,0)
    cav3[Height*20-1,:] = (0,255,0)
    cav3[:,Width*20-1] = (0,255,0)


    # combine the three cavases into one horizontal canvas
    canvas = np.hstack((cav1,cav2,cav3))

    cv2.imshow("canvas_all",canvas)
    cv2.waitKey(1)

    return canvas


# calculate fitness
def fitness(wave_):

    # calculate second order dirivative along x axis and y axis
    d1x = np.zeros((Height,Width), np.float32)
    d1y = np.zeros((Height,Width), np.float32)
    d2x = np.zeros((Height,Width), np.float32)
    d2y = np.zeros((Height,Width), np.float32)

    # get d1
    for j in range(Height):
        for i in range(Width):
            i_0 = i-1
            j_0 = j-1
            if i_0 >= 0:
                w1 = wave_[j*Width+i][0][0]
                w2 = wave_[j*Width+i_0][0][0]
                if w1 in range(6,10):
                    w1 -= 5
                elif w1 in range(10,14):
                    w1 -= 8.5
                if w2 in range(6,10):
                    w2 -= 5
                elif w2 in range(10,14):
                    w2 -= 8.5
                d1x[j,i_0] = w1 - w2
            if j_0 >= 0:
                w1 = wave_[j_0*Width+i][0][0]
                w2 = wave_[j*Width+i][0][0]
                if w1 in range(6,10):
                    w1 -= 5
                elif w1 in range(10,14):
                    w1 -= 8.5
                if w2 in range(6,10):
                    w2 -= 5
                elif w2 in range(10,14):
                    w2 -= 8.5
                d1y[j_0,i] = w1 - w2
    # get d2
    for j in range(Height):
        for i in range(Width):
            i_0 = i-1
            j_0 = j-1
            if i_0 >= 0:
                d2x[j,i_0] = d1x[j,i] - d1x[j,i_0]
            if j_0 >= 0:
                d2y[j_0,i] = d1y[j,i] - d1y[j_0,i]

    # compute sum of absolute value of d2
    smoothness = 0
    for j in range(Height):
        for i in range(Width):
            smoothness += abs(d2x[j,i]) + abs(d2y[j,i])

    dir_x = [0.0001,0.0001]
    dir_y = [0.0001,0.0001]
    for j in range(Height-1):
        for i in range(Width-1):
            if d1x[j,i] > 0:
                dir_x[0] += 1
            elif d1x[j,i] < 0:
                dir_x[1] += 1
            # else:
            #     dir_x[2] += 1
            if d1y[j,i] > 0:
                dir_y[0] += 1
            elif d1y[j,i] < 0:
                dir_y[1] += 1
            # else:
            #     dir_y[2] += 1

    # compute shannon entropy of dir_x and dir_y
    entropy_x = 0
    entropy_y = 0
    sum_x = 0.001
    sum_y = 0.001
    for i in range(len(dir_x)):
        sum_x += dir_x[i]
        sum_y += dir_y[i]
    for i in range(len(dir_x)):
        if sum_x != 0:
            entropy_x += -dir_x[i]/sum_x * math.log(dir_x[i]/sum_x)
        if sum_y != 0:
            entropy_y += -dir_y[i]/sum_y * math.log(dir_y[i]/sum_y)

    # compute balanceness of element distribution
    balance = 0.0001
    candidates_n = candidates_weight.copy()
    for i in range(len(candidates_n)):
        candidates_n[i] = 0
    for j in range(Height):
        for i in range(Width):
            w1 = wave_[j*Width+i][0][0]
            if w1 in range(6,10):
                candidates_n[w1-4-1] += 0.5
                candidates_n[w1-5-1] += 0.5
            elif w1 in range(10,14):
                candidates_n[w1-8-1] += 0.5
                candidates_n[w1-9-1] += 0.5
            else:
                candidates_n[w1-1] += 1
                
    # compute average of candidates_n
    s_ = sum(candidates_n)
    for i in range(len(candidates_n)):
        candidates_n[i]/=s_
    avg_ = sum(candidates_n)/len(candidates_n)
    for i in range(len(candidates_n)):
        balance += abs(candidates_n[i]-avg_)

    # count concave corner
    concave_corners = 0
    for j in range(Height):
        for i in range(Width):
            w = wave_[j*Width+i][0][0]
            if w in range(1,6):
                d0 = (i+1,j)
                d1 = (i,j+1)
                d2 = (i-1,j)
                d3 = (i,j-1)
                w0 = 0
                w1 = 0
                w2 = 0
                w3 = 0
                if d0[0] < Width:
                    w0 = wave_[d0[1]*Width+d0[0]][0][0]
                    if not w0 in range(1,6):
                        w0 = 0
                if d1[1] < Height:
                    w1 = wave_[d1[1]*Width+d1[0]][0][0]
                    if not w1 in range(1,6):
                        w1 = 0
                if d2[0] >= 0:
                    w2 = wave_[d2[1]*Width+d2[0]][0][0]
                    if not w2 in range(1,6):
                        w2 = 0
                if d3[1] >= 0:
                    w3 = wave_[d3[1]*Width+d3[0]][0][0]
                    if not w3 in range(1,6):
                        w3 = 0
                if w0 > w and w1 >w and w2 <= w and w3 <= w:
                    concave_corners += 1
                if w1 > w and w2 >w and w3 <= w and w0 <= w:
                    concave_corners += 1
                if w2 > w and w3 >w and w0 <= w and w1 <= w:
                    concave_corners += 1
                if w3 > w and w0 >w and w1 <= w and w2 <= w:
                    concave_corners += 1
                

    fitness_ =  7 * (abs(entropy_x) + abs(entropy_y)) + smoothness - 10 * balance + 2.2 * concave_corners # - 5 * abs(min_-max_)

    # print fitness
    print("fitness:",7 * (abs(entropy_x) + abs(entropy_y)),smoothness,10 * balance,concave_corners)

    # add a small purturbation to fitness
    fitness_ += np.random.rand() * 1

    return fitness_

# initialize
def init():

    build_entropy_lut()

    # rotate candidates
    oriented_candidates.clear()
    oriented_candidates_weight.clear()
    wave_oriented.clear()
    oriented_weight_stack.clear()
    oriented_weight_dict.clear()
    entropy.clear()

    for i in range(len(candidates)):
        candidate_0 = (candidates[i],0)
        candidate_90 = rotate_clockwise(candidate_0)
        candidate_180 = rotate_clockwise(candidate_90)
        candidate_270 = rotate_clockwise(candidate_180)
        oriented_candidates.append(candidate_0)
        oriented_candidates.append(candidate_90)
        oriented_candidates.append(candidate_180)
        oriented_candidates.append(candidate_270)
        oriented_candidates_weight.append(candidates_weight[i])
        oriented_candidates_weight.append(candidates_weight[i])
        oriented_candidates_weight.append(candidates_weight[i])
        oriented_candidates_weight.append(candidates_weight[i])

        oriented_candidates_weight_dict[candidate_0] = candidates_weight[i]
        oriented_candidates_weight_dict[candidate_90] = candidates_weight[i]
        oriented_candidates_weight_dict[candidate_180] = candidates_weight[i]
        oriented_candidates_weight_dict[candidate_270] = candidates_weight[i]

    for j in range(Height):
        for i in range(Width):
            oriented_weight_stack.append(oriented_candidates_weight.copy())
            wave_oriented.append(oriented_candidates.copy())
            entropy.append(1)
            oriented_weight_dict.append(oriented_candidates_weight_dict.copy())

    # print(oriented_weight_dict[12][2])
    # input()

# main function
def main():

    # init
    wave_old = []
    wave_initial = []

    global rules
    global rules_h
    global rules_v
    
    # seed the random number generator
    np.random.seed(int(time.time()))
    
    init()
    rules_ = symmetry_generator(rules)
    rules = rules_
    rules_h, rules_v = separate_rules(rules)

    step_ = 0
    new_weight = Width * Height * 2
    episode_ = 500
    f_min = 1000000
    img_id_ = 0

    energy = []
 
    while True:

        time1 = 0
        time2 = 0
        time3 = 0
        time4 = 0
        time5 = 0
        time6 = 0
        time7 = 0
                
        clear_screen()
        
        start = time.time()

        # step_ += 1
        
        # initialize the 2D-wave
        wave_oriented.clear()
        entropy.clear()
        for j in range(Height):
            for i in range(Width):
                # oriented_weight_stack.append(oriented_candidates_weight.copy())
                wave_oriented.append(oriented_candidates.copy())
                entropy.append(1)

        if len(wave_old) != 0:
            # alter the weight_stack
            for j in range(Height):
                for i in range(Width):
                    c = wave_old[j*Width+i][0]
                    oriented_weight_dict[j*Width+i] = oriented_candidates_weight_dict.copy()
                    oriented_weight_dict[j*Width+i][c] = new_weight

        time1 = time.time() - start
        start = time.time()
        
        update_wave_entropy()
        old_wave = []
        # update_wave_entropy_c(entropy,old_wave,wave_oriented, oriented_weight_dict)

        steps = 0

        time2 = time.time() - start
        start = time.time()

        while not is_collapsed(wave_oriented):

            steps += 1

            start = time.time()
            # 1. select the minimum entropy wavelet 
            collapse_element = find_minimum_entropy_position()
            time3 += time.time() - start
            start = time.time()

            # 2. collapse the minimum entropy wavelet
            collapse_at_position(collapse_element)
            time4 += time.time() - start
            start = time.time()

            # render2(wave_oriented)

            # 3. propagate the wave function
            propagate2(rules_h,rules_v,wave_oriented,Width,Height)
            time5 += time.time() - start
            start = time.time()

            update_wave_entropy()
            time6 += time.time() - start
            
            # print(entropy)
            # input()
            # update_wave_entropy_c(entropy,old_wave,wave_oriented, oriented_weight_dict)
            # propagate()

            # make a deep copy of wave_oriented
            # old_wave = copy.deepcopy(wave_oriented)
            
            # clear_screen()
            # render(wave_oriented)
            # for i in range(Height):
            #     for j in range(Width):
            #         entropy_str = str(round(entropy[i*Width+j],2))
            #         print(entropy_str,end=" ")
            #     print("")
            # input()

        # time3 = time.time() - start
        # print times
        print("step:",step_,"time:",time1,"time2:",time2,"time3:",time3,"time4:",time4,"time5:",time5,"time6:",time6)

        # cv2.waitKey(0)

        # 4. check for contradiction
        if contradiction_check():

            step_ += 1

            f_ = fitness(wave_oriented)
            
            write_ = False

            if len(wave_old) == 0:
                wave_old = wave_oriented.copy()
                wave_initial = wave_oriented.copy()

            if f_ < f_min:
                wave_old = wave_oriented.copy()
                f_min = f_

            energy.append(f_min)

            print("fitness:",f_)
            print("fitness_min:",f_min)
            # render(wave_oriented)
            # render(wave_old,"wave_old",write_,img_id_)
            # render(wave_initial,"wave_initial")

            cav = super_render(wave_oriented,wave_old,wave_initial)

            # with open(str(step_)+"child" + ".json", 'w') as outfile:
            #     json.dump(wave_oriented, outfile)
            # with open(str(step_)+"min" + ".json", 'w') as outfile:
            #     json.dump(wave_old, outfile)
            # with open(str(step_)+"initial" + ".json", 'w') as outfile:
            #     json.dump(wave_initial, outfile)

            # # save cav to file
            # cv2.imwrite("cav_"+str(step_)+".png",cav)

            if step_ == episode_ -1:
                img_id_ += 1
                # render current wave_old to file
                write_ = True
                new_weight = Width * Height * 0.9
            elif step_ >= episode_:
                wave_old = wave_oriented.copy()
                wave_initial = wave_oriented.copy()
                step_ = 0
                new_weight = Width * Height
                f_min = 1000000

                plt.plot(energy)
                plt.show()
                cv2.waitKey(0)
                exit()

            # wave_old = wave_oriented.copy()

            print("-----------------------------------------------------")
            print("No contradiction!")
            
        else:
            print("Contradiction!")

        print("step: ", step_)

# main entry
if __name__ == "__main__":
    main()
